package cn.dlc.com.tree;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;

/**
 * @author 邓立川
 * @date 2021/3/27 15:29
 * @desc 红黑树的实现
 */
@Data
public class RedBlackTree<Key extends Comparable<Key> , Value> {


    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    private class Node{
        private Key key;
        private Value value;
        private Node left;
        private Node right;
        private boolean color;

    }

    private Node root;
    private int num;
    private static  final boolean RED = true;
    private static  final boolean BLACK = false;

    public int size(){
        return this.num;
    }

    private boolean isRed(Node x){
        if(x == null) {
            return false;
        }

        // 判断是否为红色
        return x.isColor() == RED;
    }

    // 左旋
    private Node rotateLeft(Node h) {
        // 获取它的右子结点，。因为我们左旋的就是右子结点
        Node rightNode = h.getRight();

        // 将右子结点的左子节点作为父节点的右节点
        h.setRight(rightNode.getLeft());
        // 将h节点设置为右子结点的左子节点
        rightNode.setLeft(h);
        // 让右子结点的颜色属性等于h节点的颜色属性
        rightNode.setColor(h.isColor());
        // 当前节点的color设置为红色
        h.setColor(RED);

        return rightNode;
    }

    // 右旋
    private Node rotateRight(Node h) {
        // 获取它的左子节点，因为右旋的就是左子节点
        Node leftNode = h.getLeft();
        // 将左子节点的右子结点设置为当前节点的左子节点
        h.setLeft(leftNode.getRight());
        // 设置当前节点为左子节点的右子结点
        leftNode.setRight(h);
        // 左子节点的color设置为当前节点的color
        leftNode.setColor(h.isColor());
        // 当前节点的颜色设置为红色
        h.setColor(RED);

        return leftNode;
    }

    // 颜色反转
    private void flipColors(Node h) {
        // 自己变为红色
        h.setColor(RED);
        // 左右子节点变成黑色
        h.getLeft().setColor(BLACK);
        h.getRight().setColor(BLACK);
    }

    // 插入方法
    public void put(Key key, Value value) {
        // 添加元素
        this.root = put(this.root, key,value);
        // 将根节点的颜色设置为黑色
        root.setColor(BLACK);
    }

    // 重载添加方法
    private Node put(Node node, Key key, Value value) {

        if(node == null) {
            // 那么添加的就是第一个节点
            this.num++;
            return new Node(key, value, null, null, BLACK);
        }
        // node不为null那么就是有节点，进行下面的逻辑
        // 比较插入的key和当前节点的大小
        int result = key.compareTo(node.getKey());
        if(result >0) {
            // 当前节点大,设置为右边的节点，继续遍历
            node.setRight(put(node.getRight(), key, value));

        } else if(result == 0) {
            // =当前节点，重新设置值
            node.setValue(value);
        } else {
            // 小于当前节点
            node.setLeft(put(node.getLeft(), key, value));
        }

        // 插入之后需要进行左旋或者右旋或者颜色反转的调整
        // 如果node的左子节点为黑色，右子结点为红色，左旋
        if(isRed(node.getRight()) && !isRed(node.getLeft())) {
            node = rotateLeft(node);
        }

        // 如果左子节点和左子节点的左子节点为红色，进行右旋
        if(isRed(node.getLeft()) && isRed(node.getLeft().getLeft())) {
            node = rotateRight(node);
        }

        // 左右子节点都是红色
        if(isRed(node.getLeft()) && isRed(node.getRight())) {
            flipColors(node);
        }
        return node;
    }

    public Value get(Key key) {
        return get(root, key);
    }

    // 重载获取方法
    private Value get(Node node, Key key) {

        if(node == null) {
            return null;
        }

        int result = key.compareTo(node.getKey());
        if(result > 0) {
            // 大于
            return get(node.getRight(), key);
        } else if(result < 0) {
            return get(node.getLeft(), key);
        } else {
            return node.getValue();
        }
    }
}
